# coding=utf-8
# Copyright 2018 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Convert OpenAI GPT checkpoint."""

from __future__ import absolute_import, division, print_function
from fairseq import options, tasks
import torch
import logging
import os

logger = logging.getLogger(__name__)


def name_to_variable(model, name):
    r"""Find the corresponding variable given the specified name.
    """
    pointer = model
    name = name.split(".")
    for m_name in name:
        if m_name.isdigit():
            num = int(m_name)
            pointer = pointer[num]  # type: ignore
        else:
            pointer = getattr(pointer, m_name)
    return pointer


def load_tf_weights_in_gpt2(model, gpt2_checkpoint_path):
    """ Load tf checkpoints in a pytorch model
    """
    try:
        import re
        import numpy as np
        import tensorflow as tf
    except ImportError:
        logger.error("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions.")
        raise

    global_tensor_map = {
        "model/wte": "decoder.embed_tokens.weight",
        "model/wpe": "decoder.embed_positions.weight",
        "model/ln_f/b": "decoder.layer_norm.bias",
        "model/ln_f/g": "decoder.layer_norm.weight",
    }
    layer_tensor_map = {
        "ln_1/b": "decoder.layers.{}.self_attn_layer_norm.bias",
        "ln_1/g": "decoder.layers.{}.self_attn_layer_norm.weight",
        "ln_2/b": "decoder.layers.{}.final_layer_norm.bias",
        "ln_2/g": "decoder.layers.{}.final_layer_norm.weight",
        "mlp/c_fc/b": "decoder.layers.{}.fc1.bias",
        "mlp/c_proj/b": "decoder.layers.{}.fc2.bias",
        "attn/c_attn/b": "decoder.layers.{}.self_attn.in_proj_bias",
        "attn/c_proj/b": "decoder.layers.{}.self_attn.out_proj.bias",
    }
    layer_transpose_map = {
        "mlp/c_fc/w": "decoder.layers.{}.fc1.weight",
        "mlp/c_proj/w": "decoder.layers.{}.fc2.weight",
        "attn/c_attn/w": "decoder.layers.{}.self_attn.in_proj_weight",
        "attn/c_proj/w": "decoder.layers.{}.self_attn.out_proj.weight",
    }

    tf_path = os.path.abspath(os.path.join(gpt2_checkpoint_path, 'model.ckpt'))
    logger.info("Converting TensorFlow checkpoint from {}".format(tf_path))
    # Load weights from TF model
    init_vars = tf.train.list_variables(tf_path)
    names = []
    arrays = []
    for name, shape in init_vars:
        logger.info("Loading TF weight {} with shape {}".format(name, shape))
        array = tf.train.load_variable(tf_path, name)
        names.append(name)
        arrays.append(array.squeeze())

    for name, array in zip(names, arrays):
        pointer = None
        if name in global_tensor_map:
            v_name = global_tensor_map[name]
            pointer = name_to_variable(model, v_name)
        else:
            name_tmp = name.split("/")
            layer_no = name_tmp[1][1:]
            name = "/".join(name_tmp[2:])

            if name in layer_tensor_map:
                v_name = layer_tensor_map[name].format(layer_no)
                pointer = name_to_variable(model, v_name)
            elif name in layer_transpose_map:
                v_name = layer_transpose_map[name].format(layer_no)
                pointer = name_to_variable(model, v_name)
                array = np.transpose(array)
            else:
                logger.info("Unkown weight {}".format(name))
        try:
            assert pointer.shape == array.shape
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
        logger.info("Initialize PyTorch weight {}".format(name))
        pointer.data = torch.from_numpy(array)
    return model


def convert_gpt2_checkpoint_to_pytorch(gpt2_checkpoint_path):
    # Setup task, e.g., translation, language modeling, etc.
    # Print args
    print(args)
    task = tasks.setup_task(args)
    model = task.build_model(args)

    # Load weights from numpy
    load_tf_weights_in_gpt2(model, gpt2_checkpoint_path)

    # Save pytorch-model
    pytorch_weights_dump_path = os.path.join(gpt2_checkpoint_path, 'gpt2model.pytorch')
    print("Save PyTorch model to {}".format(pytorch_weights_dump_path))
    torch.save(model.state_dict(), pytorch_weights_dump_path)


def main(_args):
    convert_gpt2_checkpoint_to_pytorch(_args.gpt2_checkpoint_path)


if __name__ == "__main__":
    parser = options.get_training_parser()
    parser.add_argument("--gpt2_checkpoint_path",
                        default=None,
                        type=str,
                        required=True,
                        help="Path the TensorFlow checkpoint path.")
    args = options.parse_args_and_arch(parser)
    main(args)

